# Importing the required libraries
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
# PyTorch Library: PyTorch Team. (2025). torch module documentation. Retrieved February 25, 2025, from https://pytorch.org/docs/stable/index.html
# Torchvision: PyTorch Team. (2025). torchvision module documentation. Retrieved February 25, 2025, from https://pytorch.org/vision/stable/index.html
# Matplotlib: Hunter, J. D. (2007). Matplotlib: A 2D Graphics Environment. Computing in Science & Engineering, 9(3), 90-95. Retrieved February 25, 2025, from https://matplotlib.org/stable/index.html
# NumPy: Harris, C. R., et al. (2020). Array programming with NumPy. Nature, 585(7825), 357–362. Retrieved February 25, 2025, from https://numpy.org/doc/stable/

# Step 1: Define Device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f'Using device: {device}')
# [Source: https://pytorch.org/docs/stable/notes/cuda.html]
# https://pytorch.org/tutorials/beginner/basics/buildmodel_tutorial.html#model-layers

Using device: cpu

# Step 2: Load and Normalise Data
# Download and dataloader data from the CIFAR10 as shown in https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html
# We use transforms to change attributes of the data to make it appropriate for the pipeline.
data_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),  # Randomly flips images horizontally to introduce variation
    transforms.RandomRotation(10),  # Rotates images by a small angle to enhance model robustness
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),  # Adjusts brightness/contrast
    transforms.ToTensor(),  # Converts images to tensors for PyTorch compatibility
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # Normalizes pixel values to improve learning stability
])
# Loading CIFAR-10 dataset [Source: https://pytorch.org/vision/stable/datasets.html#torchvision.datasets.CIFAR10]
train_dataset = torchvision.datasets.CIFAR10(root='./data', train=True, transform=data_transform, download=True)
test_dataset = torchvision.datasets.CIFAR10(root='./data', train=False, transform=data_transform, download=True)

# DataLoader allows efficient batch loading [Source: https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader]
train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False)

# https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html
# PyTorch Team. (2025). torch.cuda documentation. Retrieved February 25, 2025, from https://pytorch.org/docs/stable/cuda.html

Files already downloaded and verified
Files already downloaded and verified

# Step 3: Define CNN Model with Batch Normalisation and Dropout
class CNN(nn.Module):  # [Source: PyTorch Official Examples]
    def __init__(self):
        super(CNN, self).__init__()

        # First convolutional layer extracts low-level features [Source: https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html]
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)  # Increased Filters
        self.bn1 = nn.BatchNorm2d(64)

        # Second convolutional layer extracts mid-level features.
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(128)

        # Third convolutional layer extracts deeper patterns.
        self.conv3 = nn.Conv2d(128, 256, kernel_size=3, padding=1)  # Added Extra Layer
        self.bn3 = nn.BatchNorm2d(256)

        # Max pooling reduces spatial dimensions while retaining important features. [Source: https://pytorch.org/docs/stable/generated/torch.nn.MaxPool2d.html]
        self.pool = nn.MaxPool2d(2, 2)

        # Dropout prevents overfitting [Source: https://pytorch.org/docs/stable/generated/torch.nn.Dropout.html]
        self.dropout = nn.Dropout(0.4)  # Adjusted Dropout

        # Fully connected layers for classification.
        self.fc1 = nn.Linear(256 * 4 * 4, 256)  # Adjusted Fully Connected Layer
        self.fc2 = nn.Linear(256, 10)

    # Define the forward pass of the neural network
    def forward(self, x):
        x = self.pool(torch.relu(self.bn1(self.conv1(x)))) # Apply first convolution and ReLU activation
        x = self.pool(torch.relu(self.bn2(self.conv2(x)))) # Apply second convolution and ReLU activation
        x = self.pool(torch.relu(self.bn3(self.conv3(x)))) # Apply third convolution and ReLU activation
        x = torch.flatten(x, 1) # Flattens feature maps into a vector for FC layers [Source: https://pytorch.org/docs/stable/generated/torch.flatten.html]
        x = torch.relu(self.fc1(x))
        x = self.dropout(x) # Applies dropout for regularisation [Source: https://pytorch.org/docs/stable/generated/torch.nn.Dropout.html]
        x = self.fc2(x) # Final classification layer [Source: https://pytorch.org/docs/stable/generated/torch.nn.Linear.html]
        return x

# Create an instance of the model and move it to the selected device.
model = CNN().to(device)

# https://medium.com/@myringoleMLGOD/simple-convolutional-neural-network-cnn-for-dummies-in-pytorch-a-step-by-step-guide-6f4109f6df80
# PyTorch Team. (2025). torchvision.transforms documentation. Retrieved February 25, 2025, from https://pytorch.org/vision/stable/transforms.html
# Normalisation values for CIFAR-10: Krizhevsky, A. (2009). Learning Multiple Layers of Features from Tiny Images. University of Toronto.

# Step 4: Define Loss, Optimiser, and Scheduler

criterion = nn.CrossEntropyLoss() # Cross-entropy loss for multi-class classification
optimiser = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)  # Stochastic Gradient Descent with momentum 
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimiser, T_max=15)  # Cosine Annealing Learning Rate 

# [Source: https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html]
# [Source: https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.CosineAnnealingLR.html]
# [Source: https://pytorch.org/docs/stable/generated/torch.optim.SGD.html]

# Step 5: Training Loop
def train_model(model, train_loader, test_loader, criterion, optimiser, scheduler, num_epochs=15):
    model.train()
    train_losses, val_losses = [], [] # Arrays used to store losses 
    train_accuracies, val_accuracies = [], [] # Arrays used to store calculated accuracies
    epochs = []

    
    for epoch in range(num_epochs):
        running_loss, correct_train, total_train = 0.0, 0, 0 # Values that store the number of 
        for images, labels in train_loader: # For each image and classification within the set of training data
            images, labels = images.to(device), labels.to(device) #Send the input to the device
            optimiser.zero_grad() # Sets the grads to zero, increasing performance
            outputs = model(images)
            loss = criterion(outputs, labels) # Calculates how far the predicted values are from the true labels
            loss.backward() # Backpropagation to compute gradients 
            optimiser.step()
            running_loss += loss.item()
            _, predicted = torch.max(outputs, 1)
            total_train += labels.size(0)
            correct_train += (predicted == labels).sum().item()
        
        train_loss = running_loss / len(train_loader)
        train_acc = 100 * correct_train / total_train
        train_losses.append(train_loss)
        train_accuracies.append(train_acc)
        epochs.append(epoch + 1)
        
        model.eval() # Sets the model to evaluation mode for validation
        val_loss, correct_val, total_val = 0.0, 0, 0
        with torch.no_grad(): # Disable gradient computation for validation
            for images, labels in test_loader:
                images, labels = images.to(device), labels.to(device)
                outputs = model(images)
                loss = criterion(outputs, labels) # Compute validation loss
                val_loss += loss.item()
                _, predicted = torch.max(outputs, 1)
                total_val += labels.size(0)
                correct_val += (predicted == labels).sum().item()
        
        val_loss /= len(test_loader)
        val_acc = 100 * correct_val / total_val
        val_losses.append(val_loss)
        val_accuracies.append(val_acc)
        scheduler.step() # Adjust learning rate
        
        print(f'Epoch {epoch+1}/{num_epochs} | Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}% | Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.2f}%')
        model.train() # Switch back to training mode
        
    # Plot Training Results
    plt.figure(figsize=(12, 6))
    plt.subplot(1, 2, 1)
    plt.plot(epochs, train_losses, label='Train Loss')
    plt.plot(epochs, val_losses, label='Validation Loss')
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    plt.title('Loss Over Epochs')
    plt.legend()
    
    plt.subplot(1, 2, 2)
    plt.plot(epochs, train_accuracies, label='Train Accuracy')
    plt.plot(epochs, val_accuracies, label='Validation Accuracy')
    plt.xlabel('Epochs')
    plt.ylabel('Accuracy (%)')
    plt.title('Accuracy Over Epochs')
    plt.legend()
    plt.show()
    
    return train_losses, train_accuracies, val_losses, val_accuracies

train_model(model, train_loader, test_loader, criterion, optimiser, scheduler, num_epochs=15) # Run the model
# https://pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html
# https://pytorch.org/docs/stable/tensors.html#torch.Tensor.to
# https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.forward
# https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html
# https://pytorch.org/docs/stable/autograd.html
# https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.CosineAnnealingLR.html
# https://pytorch.org/docs/stable/generated/torch.optim.Optimizer.step.html
# https://pytorch.org/docs/stable/generated/torch.max.html

Epoch 1/15 | Train Loss: 1.5120, Train Acc: 44.61% | Val Loss: 1.3842, Val Acc: 50.04%
Epoch 2/15 | Train Loss: 1.1908, Train Acc: 57.72% | Val Loss: 1.1042, Val Acc: 61.23%
Epoch 3/15 | Train Loss: 1.0705, Train Acc: 62.56% | Val Loss: 1.0589, Val Acc: 62.85%
Epoch 4/15 | Train Loss: 0.9675, Train Acc: 66.00% | Val Loss: 0.8656, Val Acc: 69.88%
Epoch 5/15 | Train Loss: 0.8902, Train Acc: 68.84% | Val Loss: 0.8533, Val Acc: 69.80%
Epoch 6/15 | Train Loss: 0.8422, Train Acc: 70.63% | Val Loss: 0.7448, Val Acc: 73.83%
Epoch 7/15 | Train Loss: 0.7947, Train Acc: 72.56% | Val Loss: 0.8082, Val Acc: 72.43%
Epoch 8/15 | Train Loss: 0.7455, Train Acc: 73.94% | Val Loss: 0.7633, Val Acc: 73.54%
Epoch 9/15 | Train Loss: 0.7078, Train Acc: 75.57% | Val Loss: 0.6910, Val Acc: 76.16%
Epoch 10/15 | Train Loss: 0.6712, Train Acc: 76.87% | Val Loss: 0.6521, Val Acc: 77.27%
Epoch 11/15 | Train Loss: 0.6344, Train Acc: 77.94% | Val Loss: 0.6203, Val Acc: 78.21%
Epoch 12/15 | Train Loss: 0.6052, Train Acc: 78.98% | Val Loss: 0.6227, Val Acc: 78.06%
Epoch 13/15 | Train Loss: 0.5832, Train Acc: 79.70% | Val Loss: 0.5955, Val Acc: 78.96%
Epoch 14/15 | Train Loss: 0.5711, Train Acc: 80.47% | Val Loss: 0.5754, Val Acc: 79.94%
Epoch 15/15 | Train Loss: 0.5546, Train Acc: 80.93% | Val Loss: 0.5767, Val Acc: 79.79%

([1.512048579542838,
  1.1908392604354703,
  1.0704513658647952,
  0.9675115708195036,
  0.8902251476521992,
  0.8421515398623084,
  0.7946918207361265,
  0.7455178170710268,
  0.7078419516763419,
  0.6712166821712728,
  0.634432458435483,
  0.6052003820686389,
  0.5832140643121032,
  0.5711386279224435,
  0.5546144381203615],
 [44.612,
  57.722,
  62.56,
  65.996,
  68.836,
  70.634,
  72.562,
  73.942,
  75.566,
  76.872,
  77.938,
  78.98,
  79.696,
  80.474,
  80.926],
 [1.3841868261747723,
  1.1042084618459773,
  1.0589210270326348,
  0.8656254976610595,
  0.853308998331239,
  0.744849595087993,
  0.8082426763788054,
  0.7632524982283387,
  0.6909980739973769,
  0.6520906713189958,
  0.6203382151036323,
  0.6226527996455566,
  0.5955363478087172,
  0.5753705584550206,
  0.5767255289645135],
 [50.04,
  61.23,
  62.85,
  69.88,
  69.8,
  73.83,
  72.43,
  73.54,
  76.16,
  77.27,
  78.21,
  78.06,
  78.96,
  79.94,
  79.79])

# Step 6: Visualising Predictions
def visualise_predictions(model, test_loader):
    classes = ('airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
    model.eval() # Set model to evaluation mode
    images, labels = next(iter(test_loader)) # Get a batch of test images
    images, labels = images.to(device), labels.to(device)
    outputs = model(images) # Forward pass to get predictions
    _, preds = torch.max(outputs, 1) # Get predicted class indices
    
    fig, axes = plt.subplots(3, 3, figsize=(8, 8)) # Create a 3x3 grid for images
    axes = axes.flatten()
    for i in range(9):
        img = images[i].cpu().numpy().transpose((1, 2, 0)) # Convert image to NumPy and adjust dimensions
        img = (img * 0.5) + 0.5  # Unnormalise image for display
        axes[i].imshow(img) # Display image
        axes[i].set_title(f'True: {classes[labels[i]]}\nPred: {classes[preds[i]]}')  # Set title with true and predicted labels
        axes[i].axis('off') # Hide axes
    plt.show() # Display the figure

visualise_predictions(model, test_loader)  # Call visualisation function

# https://pytorch.org/tutorials/intermediate/tensorboard_tutorial.html
# https://matplotlib.org/stable/tutorials/introductory/pyplot.html

# Step 7: Hyperparameter Effects on Model Performance
def test_hyperparameters(learning_rates, batch_sizes, epochs=10):
    """Tests different hyperparameter settings and compares performance."""
    results = []
    for lr in learning_rates:
        for batch_size in batch_sizes:
            print(f"Training with Learning Rate: {lr}, Batch Size: {batch_size}")
            
            # Reload dataset with new batch size
            train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
            test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
            
            # Reinitialise model, optimiser, and scheduler
            model = CNN().to(device)
            optimiser = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=5e-4)
            scheduler = optim.lr_scheduler.CosineAnnealingLR(optimiser, T_max=epochs)
            
            # Train model
            for epoch in range(epochs):
                train_losses, train_accuracies, val_losses, val_accuracies = train_model(
                    model, train_loader, test_loader, criterion, optimiser, scheduler, num_epochs=1
                )
                
                # Print out results per epoch
                print(f"Epoch {epoch+1}/{epochs} | LR: {lr}, Batch Size: {batch_size} | Train Loss: {train_losses[-1]:.4f}, Train Accuracy: {train_accuracies[-1]:.2f}% | Validation Loss: {val_losses[-1]:.4f}, Validation Accuracy: {val_accuracies[-1]:.2f}%")
                print("------------------------------------------------------")
                
            # Store final results
            results.append({
                "learning_rate": lr,
                "batch_size": batch_size,
                "train_accuracy": train_accuracies[-1],
                "val_accuracy": val_accuracies[-1]
            })
    
    # Plot how hyperparameters affect accuracy
    fig, ax = plt.subplots(figsize=(8, 6))
    for res in results:
        ax.scatter(res["learning_rate"], res["val_accuracy"], label=f"Batch Size {res['batch_size']}", s=100)
    
    ax.set_xlabel("Learning Rate")
    ax.set_ylabel("Validation Accuracy (%)")
    ax.set_title("Impact of Learning Rate and Batch Size on Model Accuracy")
    ax.set_xscale("log")  # Log scale for better visualization of learning rate differences
    ax.set_ylim(0, 100)  # Ensure accuracy is displayed between 0-100%
    ax.legend()
    plt.show()

# Run hyperparameter tests
test_hyperparameters(learning_rates=[0.1, 0.01, 0.001], batch_sizes=[32, 64, 128], epochs=10)

# https://machinelearningmastery.com/grid-search-hyperparameters-deep-learning-models-python-keras/
# https://towardsdatascience.com/understanding-learning-rates-and-how-it-impacts-your-deep-learning-model-1cd6b49197c3
# https://matplotlib.org/stable/tutorials/introductory/pyplot.html

Training with Learning Rate: 0.1, Batch Size: 32
Epoch 1/1 | Train Loss: 2.2083, Train Acc: 17.16% | Val Loss: 1.9620, Val Acc: 26.60%

Epoch 1/10 | LR: 0.1, Batch Size: 32 | Train Loss: 2.2083, Train Accuracy: 17.16% | Validation Loss: 1.9620, Validation Accuracy: 26.60%
------------------------------------------------------
Epoch 1/1 | Train Loss: 1.9482, Train Acc: 26.12% | Val Loss: 1.8579, Val Acc: 28.70%

Epoch 2/10 | LR: 0.1, Batch Size: 32 | Train Loss: 1.9482, Train Accuracy: 26.12% | Validation Loss: 1.8579, Validation Accuracy: 28.70%
------------------------------------------------------
Epoch 1/1 | Train Loss: 1.8886, Train Acc: 29.68% | Val Loss: 1.9356, Val Acc: 29.13%

Epoch 3/10 | LR: 0.1, Batch Size: 32 | Train Loss: 1.8886, Train Accuracy: 29.68% | Validation Loss: 1.9356, Validation Accuracy: 29.13%
------------------------------------------------------
Epoch 1/1 | Train Loss: 1.8107, Train Acc: 33.11% | Val Loss: 1.5736, Val Acc: 40.82%

Epoch 4/10 | LR: 0.1, Batch Size: 32 | Train Loss: 1.8107, Train Accuracy: 33.11% | Validation Loss: 1.5736, Validation Accuracy: 40.82%
------------------------------------------------------
Epoch 1/1 | Train Loss: 1.7205, Train Acc: 36.64% | Val Loss: 1.6229, Val Acc: 43.41%

Epoch 5/10 | LR: 0.1, Batch Size: 32 | Train Loss: 1.7205, Train Accuracy: 36.64% | Validation Loss: 1.6229, Validation Accuracy: 43.41%
------------------------------------------------------
Epoch 1/1 | Train Loss: 1.6158, Train Acc: 40.58% | Val Loss: 1.5459, Val Acc: 43.22%

Epoch 6/10 | LR: 0.1, Batch Size: 32 | Train Loss: 1.6158, Train Accuracy: 40.58% | Validation Loss: 1.5459, Validation Accuracy: 43.22%
------------------------------------------------------
Epoch 1/1 | Train Loss: 1.5247, Train Acc: 44.28% | Val Loss: 1.5468, Val Acc: 44.28%

Epoch 7/10 | LR: 0.1, Batch Size: 32 | Train Loss: 1.5247, Train Accuracy: 44.28% | Validation Loss: 1.5468, Validation Accuracy: 44.28%
------------------------------------------------------
Epoch 1/1 | Train Loss: 1.4231, Train Acc: 48.26% | Val Loss: 1.2664, Val Acc: 54.40%

Epoch 8/10 | LR: 0.1, Batch Size: 32 | Train Loss: 1.4231, Train Accuracy: 48.26% | Validation Loss: 1.2664, Validation Accuracy: 54.40%
------------------------------------------------------
Epoch 1/1 | Train Loss: 1.3183, Train Acc: 52.49% | Val Loss: 1.2189, Val Acc: 55.37%

Epoch 9/10 | LR: 0.1, Batch Size: 32 | Train Loss: 1.3183, Train Accuracy: 52.49% | Validation Loss: 1.2189, Validation Accuracy: 55.37%
------------------------------------------------------
Epoch 1/1 | Train Loss: 1.2222, Train Acc: 56.35% | Val Loss: 1.0926, Val Acc: 60.98%

RESULTS:

RESULTS OF CHANGING HYPERPARAMETERS:

How does this tutorial differ from other similar tutorials?

Sources